package bm25;//created by lidongdong
//used to search the most relevant document 1000

import java.util.List;
import java.io.*;
import java.nio.file.Paths;
import java.util.ArrayList;

import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.queryparser.classic.ParseException;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.IndexSearcher;
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.TopDocs;
import org.apache.lucene.search.similarities.BM25Similarity;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.store.IOContext;
import org.apache.lucene.store.RAMDirectory;

public class TOP1000 {

    // 获取到特定query的top1000文档
    public static List<String> gettop1000(String query, IndexSearcher searcher, QueryParser parser) throws IOException, ParseException {
        TopDocs results = searcher.search(parser.parse(query), 1000 * 8000);
        ScoreDoc[] hists = results.scoreDocs;

        List<String> docids = new ArrayList<>(1000);
        for (int i = 0; i < 1000; i ++){
            Document doc = searcher.doc(hists[i].doc);
            String docPath = doc.get("path");
            String docID = docPath.substring(Math.max(docPath.lastIndexOf("/"),
                    docPath.lastIndexOf("\\")) + 1, docPath.lastIndexOf("."));
            docids.add(docID);
        }
        return docids;
    }

    private static IndexSearcher getSearcher() throws IOException {
        String indexPath = "/home/jack/IdeaProjects/hello_world/trec.index";
        Directory directory = new RAMDirectory(FSDirectory.open(Paths.get(indexPath)), IOContext.READONCE);
        DirectoryReader directory_reader = DirectoryReader.open(directory);
        IndexSearcher searcher = new IndexSearcher(directory_reader);
        // 参数的设置 k1 = 2.6 b = 0.65
        searcher.setSimilarity(new BM25Similarity(2.6f, 0.65f));
        return searcher;
    }


    public static void main(String[] args) throws IOException, ParseException {
        String query = args[0];
        String target_filename = args[1];  //文件名
        boolean query_is_file = false;
        if (args.length >= 3) {
            query_is_file = Boolean.getBoolean(args[2]);
        }

        // 对于检索的情况, 居然在字段不在的情况下没有打印任何信息
        QueryParser parser = new QueryParser("contents", new StandardAnalyzer());
        IndexSearcher searcher = getSearcher();

        File target_file = new File(target_filename);
        // validation 对于只有一个query的情况
        if (!target_file.exists()) {
            target_file.createNewFile();
        }
        FileWriter fw = new FileWriter(target_file);

        if (query_is_file) {
            List<String> docids = gettop1000(query, searcher, parser);
            for (int i = 0; i < docids.size(); i++) {
                String docid = docids.get(i);
                fw.write(docid);
                fw.write("\n");
            }
        }else{
            // 对于以query作为文件的情况
            List<String> topics = utils.getTopics(query);
            for (String topic:topics){
                topic = DataPreprocess.remove_stemming(topic);
                // 按照顺序写 query1 : 1 - 1000行, query2: 1001 - 2000行, ...
                List<String> docids = gettop1000(topic, searcher, parser);
                for (String docid: docids) {
                    fw.write(docid);
                    fw.write("\n");
                }
            }
        }

        fw.flush();
        fw.close();
    }
}
